﻿using System;
using System.Collections.Concurrent;
using System.Runtime.CompilerServices;
using System.Threading;
using UnityEngine;

namespace JH
{
    /// <summary>
    /// 条件等待器，满足条件后继续执行
    /// </summary>
    public struct ConditionAwaiter : ICriticalNotifyCompletion
    {
        /// <summary>
        ///  await后的回调
        /// </summary>
        private Action _continuation;

        /// <summary>
        /// 需要满足的条件
        /// </summary>
        private readonly Func<bool> _condition;

        /// <summary>
        ///  取消标记
        /// </summary>
        private readonly CancellationToken _token;

        /// <summary>
        ///  是否取消
        /// </summary>
        private readonly AwaiterState _awaiterState;

        /// <summary>
        ///  取消回调
        /// </summary>
        private readonly CancellationTokenRegistration _tokenRegistration;

        /// <summary>
        ///  调度行为状态
        /// </summary>
        private readonly DispatchState _dispatchState;

        /// <summary>
        ///  INotifyCompletion接口中的完成判定
        /// </summary>
        public bool IsCompleted => CanInvoke();

        /// <summary>
        /// 初始化或重置方法（从池中取出awaiter时调用）
        /// </summary>
        public ConditionAwaiter(Func<bool> condition, CancellationToken token = default,
            DispatchState dispatchState = DispatchState.Default)
        {
            if (token.IsCancellationRequested) throw new OperationCanceledException(token);
            _condition = condition;
            _continuation = null;
            _token = token;
            var cancelState = _awaiterState = AwaiterState.Rent();
            _dispatchState = dispatchState;
            _tokenRegistration = _token.CanBeCanceled
                ? _token.Register(() => { cancelState.IsCanceled = true; })
                : default;
        }

        /// <summary>
        ///  当IsCompleted为false时，将会进入此函数，以处理复杂逻辑
        /// </summary>
        /// <param name="continuation">await后的回调</param>
        public void OnCompleted(Action continuation)
        {
            _continuation = continuation;
            ThreadDispatcher.Invoke(CanInvoke, InvokeContinuation, _dispatchState);
        }

        /// <summary>
        ///  用于await关键字的支持
        /// </summary>
        public void UnsafeOnCompleted(Action continuation)
        {
            OnCompleted(continuation);
        }

        /// <summary>
        /// 能否调用回调的判定逻辑
        /// </summary>
        public bool CanInvoke()
        {
            if (_awaiterState != null)
                return _awaiterState.IsCanceled || _condition == null || _condition.Invoke();
            return _condition == null || _condition.Invoke();
        }

        /// <summary>
        ///  调用回调
        /// </summary>
        private void InvokeContinuation()
        {
            _continuation?.Invoke();

            // 取消注册以释放资源
            _tokenRegistration.Dispose();

            // 回收取消状态（完成后自动回到池中）
            _awaiterState.Return();
        }

        /// <summary>
        ///  设置异常
        /// </summary>
        public void SetException(Exception exception)
        {
            if (exception == null) return;
            _awaiterState.Exception = exception;
        }

        public void GetResult()
        {
            //其他由状态机传递的异常
            if (_awaiterState.Exception != null) throw _awaiterState.Exception;

            // 如果被取消，则抛出取消异常
            if (_awaiterState.IsCanceled)throw new OperationCanceledException(_token);
        }

        public ConditionAwaiter GetAwaiter() => this;

        /// <summary>
        ///  由于值类型无法持有变化的状态【它的状态修改会应用到副本上】，所以使用一个类来持有状态
        /// </summary>
        private sealed class AwaiterState
        {
            public Exception Exception;
            public bool IsCanceled;
            private static readonly ConcurrentQueue<AwaiterState> Pool = new();

            public static AwaiterState Rent()
            {
                var ret = Pool.TryDequeue(out var state) ? state : new AwaiterState();
                ret.IsCanceled = false;
                ret.Exception = null;
                return ret;
            }

            public void Return()
            {
                Pool.Enqueue(this);
            }
        }
    }
}